package cn.ycc1.my.agent.service.impl;

import cn.ycc1.my.agent.service.StreamingService;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.model.StreamingResponseHandler;
import dev.langchain4j.model.openai.OpenAiStreamingChatModel;
import dev.langchain4j.model.output.Response;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Sinks;

/**
 * @author ycc
 * @date 2025/3/1
 */
@Service
public class StreamingServiceImpl implements StreamingService {
    private final OpenAiStreamingChatModel chatModel;

    public StreamingServiceImpl(OpenAiStreamingChatModel chatModel) {
        this.chatModel = chatModel;
    }

    public Flux<String> streamResponseBak(String prompt) {
        Sinks.Many<String> sink = Sinks.many().unicast().onBackpressureBuffer();

        chatModel.generate(prompt, new StreamingResponseHandler<AiMessage>() {
            @Override
            public void onNext(String token) {
                sink.tryEmitNext(token);
            }

            @Override
            public void onComplete(Response<AiMessage> response) {
                sink.tryEmitComplete();
            }

            @Override
            public void onError(Throwable error) {
                sink.tryEmitError(error);
            }
        });

        return sink.asFlux();
    }

    /**
     * 使用Flux.create方法创建流式响应
     * 错误处理增强
     * @param prompt
     * @return
     */
    public Flux<String> streamResponse(String prompt) {
        return Flux.create(sink -> {
            chatModel.generate(prompt, new StreamingResponseHandler<AiMessage>() {
                @Override
                public void onNext(String token) {
                    if (!sink.isCancelled()) {
                        sink.next(token);
                    }
                }

                @Override
                public void onComplete(Response<AiMessage> response) {
                    if (!sink.isCancelled()) {
                        sink.complete();
                    }
                }

                @Override
                public void onError(Throwable error) {
                    if (!sink.isCancelled()) {
                        sink.error(error);
                    }
                }
            });
        });
    }
}
